package gen

import (
	"cs466/findmotif/motif"
	"math"
	"math/rand"
)

// for both of these functions, compute the limit of f rather than f when asked
// for border values (0, 1-cum_p); do this by omitting 0*log(0) terms, since
// those approach 0 as their variables approach 0.
func lbGivenFirstColumnP(cum_ic float64, cum_p float64, cols_remaining float64) func(float64) float64 {
	return func(p float64) float64 {
		p_rest := 1 - cum_p - p
		ret := cum_ic
		if p != 0 {
			ret += p*math.Log2(4*p)
		}
		if p_rest != 0 {
			ret += p_rest*math.Log2(4*p_rest/cols_remaining)
		}
		return ret
	}
}

func dLbGivenFirstColumnP(cum_p float64, cols_remaining float64) func(float64) float64 {
	return func(p float64) float64 {
		return math.Log2(4*p) - math.Log2(4*(1-cum_p-p)/cols_remaining)
	}
}

// see lbGivenFirstColumnP for explanation of these conditionals
func ubGivenFirstColumnP(cum_ic float64, cum_p float64) func(float64) float64 {
	return func(p float64) float64 {
		p_rest := 1 - cum_p - p
		ret := cum_ic
		if p != 0 {
			ret += p*math.Log2(4*p)
		}
		if p_rest != 0 {
			ret += p_rest*math.Log2(4*p_rest)
		}
		return ret
	}
}

func dUbGivenFirstColumnP(cum_p float64) func(float64) float64 {
	return func(p float64) float64 {
		return math.Log2(4*p) - math.Log2(4*(1-cum_p-p))
	}
}

const tolerance = 0.0000001
func Newton(f func(float64) float64, df func(float64) float64, fmin, fmax float64, target float64, guess float64) float64 {
	cur := guess
	prev := -1.
	for n := 0; prev != cur && n < 128; n++ {
		prev = cur

		val := f(cur) - target

		deriv := df(cur)

		// even if we want to take a big step, scale it down so we stay within
		// the range the functions are defined on
		for cur - val/deriv < fmin || cur - val/deriv > fmax {
			val /= 2
		}
		
		cur -= val / deriv
	}

	if math.Abs(prev - cur) > tolerance {
		panic("Newton-Raphson method failed to converge")
	}

	return cur
}

func minByLb(cum_ic float64, cum_p float64, cols_remaining float64, target_ic float64) float64 {
	f := lbGivenFirstColumnP(cum_ic, cum_p, cols_remaining)

	// if the lower-bound curve intersects x=0 below target_ic, then no p-value
	// is small enough to prevent reaching target_ic, so 0 is our min
	if f(0) < target_ic {
		return 0
	}

	// otherwise, intersect this curve with y=target_ic to get the minimum

	// lbGivenFirstColumnP is smallest when p = (1-cum_p-p)/cols_remaining,
	// since these sum(p*log(p)) functions are smallest when all of the p's are
	// equal.  we want to start our Newton's method guess somewhere s.t. it's
	// guaranteed to find the particular solution we're looking for, so we'll
	// choose halfway between 0 and this minimum.
	// Solving the minimum equation for p, we get p =
	// (1-cum_p)/(1+cols_remaining), so our guess is half of this.
	guess := (1 - cum_p) / (2 + 2*cols_remaining)

	return Newton(f, dLbGivenFirstColumnP(cum_p, cols_remaining), 0, 1-cum_p, target_ic, guess)
}

func maxByLb(cum_ic float64, cum_p float64, cols_remaining float64, target_ic float64) float64 {
	f := lbGivenFirstColumnP(cum_ic, cum_p, cols_remaining)

	// we can't set this column's p to be > 1-cum_p, since that would make the
	// sum of probabilities greater than 1
	right_bound := 1 - cum_p

	// if this curve intersects the right x-bound below target_ic, then no
	// p-value is large enough to prevent reaching target_ic, so right_bound is
	// our max
	if f(right_bound) < target_ic {
		return right_bound
	}

	// As last time, our guess is designed to select the right solution; this
	// time, it's to the right of the minimum.
	guess := ((1-cum_p)/(1+cols_remaining) + right_bound) / 2

	return Newton(f, dLbGivenFirstColumnP(cum_p, cols_remaining), 0, 1-cum_p, target_ic, guess)
}

// right side of the interval (remember, p must be _outside_ of this one, so
// the right side is a minimum for p)
func minByUb(cum_ic float64, cum_p float64, target_ic float64) float64 {
	f := ubGivenFirstColumnP(cum_ic, cum_p)

	// again, can't set this p > 1-cum_p
	right_bound := 1 - cum_p

	// this function is smallest when p = 1 - cum_p - p, for similar reasons to
	// the above.  if target_ic is less than the minimum, it never intersects
	// the curve, so there is no value that will make IC too large (neither a
	// min nor a max)
	low_point := (1 - cum_p) / 2

	if f(low_point) > target_ic {
		// return low_point both here and as the bogus-maximum in maxByUb so we
		// get a 0-size exclusion interval; this makes it easier to sample the
		// interval difference later
		return 0.5
	}

	// start with our guess to the right of the low point to ensure we get the
	// right solution
	guess := (right_bound + low_point) / 2

	return Newton(f, dUbGivenFirstColumnP(cum_p), 0, 1-cum_p, target_ic, guess)
}

// left side of the interval (remember, p must be _outside_ this one, so the
// left side is a maximum for p)
func maxByUb(cum_ic float64, cum_p float64, target_ic float64) float64 {
	f := ubGivenFirstColumnP(cum_ic, cum_p)

	// same low point as in minByUb
	low_point := (1 - cum_p) / 2

	if f(low_point) > target_ic {
		// as in minByUb
		return low_point
	}

	// our guess selects the left solution this time
	guess := low_point / 2

	return Newton(f, dUbGivenFirstColumnP(cum_p), 0, 1-cum_p, target_ic, guess)
}

// return p such that <p, sum_p-p, 0, 0, ...> (of total length n) has IC ic.
// TODO add tests for this; it looks weird on second inspection.  use
// Wolfram|Alpha to evaluate what it should be for a few combinations, to get
// test cases
func diffPairWithIC(ic float64, n int64, sum_p float64) (p float64) {
	icpc := func(x float64) float64 {
		return x*math.Log2(float64(n)*x) + (sum_p-x)*math.Log2(float64(n)*(sum_p-x)) - ic
	}

	high := sum_p
	p = sum_p/2
	low := 0.
	prev := -1.

	for i := icpc(p); i != 0 && p != prev; p = (high + low) / 2 {
		prev = p
		i = icpc(p)
		if i > 0 {
			high = p
		} else {
			low = p
		}
	}

	return p
}

// At index i:
// - cum_ic = sum(ic[x], x, 0, i-1)
// - cum_p  = sum(p[x], x, 0, i-1)
// - cols_remaining = n - i - 1
// - target_ic = the total IC we want for the whole distribution
//
// We want to generate a p such that the remaining probability (1 - cum_p - p)
// can still be distributed such that the total IC is target_ic.
//
// Given the above values, we can compute a curve for the lower bound on the IC
// of the rest of the column, as a function of p.  this is lbGivenFirstColumnP.
//
// Intersecting the area above this curve with the line y=target_ic, we find a
// continuous range that must contain p, if IC is going to be <= target_ic
//
// Likewise, given the above, we can compute an upper bound, which is
// ubGivenFirstColumnP.
//
// Intersecting the area below _that_ curve with y=target_ic, we get a range
// which must not contain p, if IC is going to be >= target_ic
//
// By finding the difference of these two ranges (hint: it's up to two disjoint
// ranges), we now have the set of acceptable values for p to take.  Since this
// is two continuous intervals, we can sample uniformly on it by sampling
// U[0,size(left_interval)+size(right_interval)] and mapping into the actual
// intervals.
//
// Given this, we update the cumulative values and continue on the next
// probability.
//
// A few things were simplified: first, because of the nature of the two
// curves, there are varying numbers of intersections.  The lower bound can
// have 0, 1, or 2; and the upper bound can have 0 or 2.  Our method to filter
// these out is explained in the corresponding functions' comments.
//
// Second, we need derivatives of both of these curves.  Fortunately, they're
// easy enough to implement; they are dLbGivenFirstColumnP and
// dUbGivenFirstColumnP.
//
// Finally, this is biased towards A-heavy distributions, and more generally
// distributions heavier in probabilities generated earlier.  To offset this,
// since the distribution of the probabilities is not very biased, simply
// permute the column at random before returning it.
//
// Note: while most of this code is not nucleotide-specific (i.e. hard-coded
// for 4 probabilities), the lower- and upper-bound curves and their
// derivatives are currently fixed at n=4.  To convert it, add another
// parameter to each containing the total number of elements, and recompute the
// derivatives.
func RandColumn(ic float64) (ret motif.Column) {
	cum_ic := 0.
	cum_p := 0.
	dist := make(motif.Column, 4)

	// use this method for the first n-2 elements, ...
	for i := 0; i < len(dist)-2; i++ {
		min_lb := minByLb(cum_ic, cum_p, float64(4-i-1), ic)
		max_lb := maxByLb(cum_ic, cum_p, float64(4-i-1), ic)

		min_ub := minByUb(cum_ic, cum_p, ic)
		max_ub := maxByUb(cum_ic, cum_p, ic)

		total_width := (max_ub - min_lb) + (max_lb - min_ub)

		// generate our random value and adjust it to have the appropriately
		// sized range and start at the appropriate point
		//
		// we would get _actually_-uniform distributions according to the
		// Dirichlet distribution if we managed to sample this according to a
		// chopped-up Beta distribution; however, I really don't want to
		// implement the regularized incomplete Beta function just to make my
		// PWMs a little bit more uniform.  So, it's just uniform within the
		// set.
		p := min_lb + total_width*rand.Float64()

		// if the adjusted value falls outside the first interval, shift it so
		// that the second interval's part of the distribution lines up with
		// the actual second interval
		if p > max_ub {
			p += min_ub - max_ub
		}

		dist[i] = p
		cum_p += p
		cum_ic += p*math.Log2(4*p)
	}

	// then, since there is now exactly one pair of values that can make the
	// total IC equal to ic, compute it directly.
	p := diffPairWithIC(ic - cum_ic, 4, 1 - cum_p)
	dist[len(dist)-2] = p
	dist[len(dist)-1] = 1 - cum_p - p

	// finally, shuffle!
	bag := []int{0,1,2,3}
	ret = make([]float64, 4)
	for i := range ret{
		// draw from the bag: pick a random index, move it to the front, use
		// its value, and cut it off the front of the slice
		ix := rand.Int31n(int32(len(bag)))
		bag[0], bag[ix] = bag[ix], bag[0]
		ret[i] = dist[bag[0]]
		bag = bag[1:]
	}

	return
}
